import os, cv2
import sys
sys.path.append('..')
import glob
import pandas as pd
import numpy as np
from pathlib import Path
from utils.helper import read_xml, find_nodes, change_node_text, indent, write_xml
from utils.json import json_to_image
import json


def json2xml(img_path, dst_path):
    # print(json_path)
    label_path = str(Path(img_path).with_suffix('.json'))
    if not os.path.exists(label_path):
        label_path = str(Path(img_path).with_suffix('.png'))
        if not os.path.exists(label_path):
            return 
    
    if label_path.endswith('.json'):
        try:
            img_path = str(Path(label_path).with_suffix('.jpg'))
            if not  os.path.exists(img_path): return
            img = cv2.imread(img_path, 1)
            mask = json_to_image(img.shape, label_path).astype('uint8')
            shape = img.shape[:2]
        except:
            return 
    elif label_path.endswith('.png'):
        mask = cv2.imread(label_path, 0).astype('uint8')
        shape = mask.shape

    base_tree = read_xml("../datasets/base_example.xml")
    root = base_tree.getroot()
    anno_tree = read_xml("../datasets/anno_example.xml")
    folder_node = find_nodes(base_tree, "folder")
    filename_node = find_nodes(base_tree, "filename")
    path_node = find_nodes(base_tree, "path")
    width_node = find_nodes(base_tree, "size/width")
    height_node = find_nodes(base_tree, "size/height")
    depth_node = find_nodes(base_tree, "size/depth")
    change_node_text(folder_node, Path(label_path).parent.name)
    change_node_text(filename_node, Path(label_path).name)
    change_node_text(path_node, str(Path(label_path).parent))
    change_node_text(width_node, str(shape[1]))
    change_node_text(height_node, str(shape[0]))
    change_node_text(depth_node, str(3))

    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    x_list = []
    y_list = []
    radius_list = []
    for cnt in contours:
        (x,y), radius = cv2.minEnclosingCircle(cnt)
        x_list.append(x)
        y_list.append(y)
        radius_list.append(radius)
    if len(radius_list) > 0:
        index = np.argmax(radius_list)
        # defect_size_diameter = 2 * np.nanmax(radius_list)
        x, y, w, h = cv2.boundingRect(contours[index])

        x_min, x_max, y_min, y_max = x, x+w, y, y+h
        # img_boxed = img_ori[y_min:y_max, x_min:x_max, ...]

        xmin_node = find_nodes(anno_tree, "bndbox/xmin")
        ymin_node = find_nodes(anno_tree, "bndbox/ymin")
        xmax_node = find_nodes(anno_tree, "bndbox/xmax")
        ymax_node = find_nodes(anno_tree, "bndbox/ymax")
        change_node_text(xmin_node, str(x_min))
        change_node_text(ymin_node, str(y_min))
        change_node_text(xmax_node, str(x_max))
        change_node_text(ymax_node, str(y_max))
        root.append(anno_tree.getroot())
        indent(root)

        code = Path(label_path).parent.name
        xml_name = Path(label_path).with_suffix('.xml').name
        dst = os.path.join(dst_path, code, xml_name)
        os.makedirs(os.path.join(dst_path, code), exist_ok=True)
        write_xml(base_tree, dst)


def box_scale_crop(img_path, original_bbox, crop_size=224):
    img = cv2.imread(img_path, 1)
    h, w = img.shape[:2]
    if isinstance(crop_size, int):
        crop_size = (crop_size, crop_size)  # h, w

    xmin, ymin, xmax, ymax = original_bbox
    center_x = (xmax + xmin) / 2.
    center_y = (ymax + ymin) / 2.
    defect_w, defect_h = xmax-xmin, ymax-ymin

    if defect_w/defect_h > crop_size[1] / crop_size[0]:
        new_defect_h = crop_size[0] * defect_w / crop_size[1]
        new_defect_w = defect_w
    else:
        new_defect_w = crop_size[1] * defect_h / crop_size[0]
        new_defect_h = defect_h
    
    if new_defect_h < crop_size[0]:
        new_defect_h, new_defect_w = crop_size[0], crop_size[1]

    new_x_min = max(int(center_x-int(new_defect_w/2)+1), 0)
    new_y_min = max(int(center_y-int(new_defect_h/2)+1), 0)
    new_x_max = min(int(center_x+int(new_defect_w/2)+1), w)
    new_y_max = min(int(center_y+int(new_defect_h/2)+1), h)

    bbox = [new_x_min, new_y_min, new_x_max, new_y_max]

    return bbox


def box_scale_crop_v2(img_path, original_bbox, scale=0., c1=0.7, c2=0.6, crop_size=299):
    img = cv2.imread(img_path, 1)
    height, width, _ = img.shape
    if isinstance(crop_size, int):
        crop_size = (crop_size, crop_size)  # h, w

    xmin, ymin, xmax, ymax = original_bbox
    center_x = (xmax + xmin) / 2.
    center_y = (ymax + ymin) / 2.
    bbox_w, bbox_h = xmax-xmin, ymax-ymin
    
    if bbox_h > bbox_w:
        y_side_ = bbox_h
    else:
        y_side_ = bbox_w
    if c1 is not None and c2 is not None:
        if y_side_ >= c1 * width:
            y_side_ = int(c2 * y_side_)

    y_side_ = int(y_side_ * (1 + scale))

    if y_side_ > height or y_side_ > width:
        y_side_ = min(height, width)

    if y_side_ < crop_size[0] * (1 + scale):
        y_side_ = crop_size[0] * (1 + scale)

    x_side = y_side_
    y_side = y_side_
    crop_info = [int(center_y - y_side / 2), int(center_y + y_side / 2),
                    int(center_x - x_side / 2), int(center_x + x_side / 2)]
    if crop_info[0] <= 0:
        crop_info[0] = 0
        crop_info[1] = y_side
    if crop_info[1] >= height:
        crop_info[1] = height
        crop_info[0] = max(0, height - y_side)
    if crop_info[2] <= 0:
        crop_info[2] = 0
        crop_info[3] = x_side
    if crop_info[3] >= width:
        crop_info[3] = width
        crop_info[2] = max(0, width - x_side)
    
    new_bbox = [int(crop_info[2]), int(crop_info[0]), int(crop_info[3]), int(crop_info[1])]

    return new_bbox


def crop_image(img_path, original_bbox, method):
    if method == 'original_crop':
        return original_bbox  # crop_original(img_path, original_bbox)
    elif method == '224_crop':
        return box_scale_crop(img_path, original_bbox, crop_size=224)
    elif method == '299_crop':
        return box_scale_crop(img_path, original_bbox, crop_size=299)
    elif method == '224_scale_crop':
        return box_scale_crop_v2(img_path, original_bbox, scale=0., c1=0.7, c2=0.6, crop_size=224)
    elif method == '299_scale_crop':
        return box_scale_crop_v2(img_path, original_bbox, scale=0., c1=0.7, c2=0.6, crop_size=299)
    else:
        return None 


def get_original_bbox(json_path):
    code = None
    radius2 = 0
    bbox = None

    try:
        with open(json_path, 'r+', encoding="utf-8") as f:
            """ Load json file and update dictionary order """
            data = json.load(f)

        for part in data['shapes']:
            countor = np.array(part['points'], dtype=np.int32)
            countor = np.expand_dims(countor, axis=1)
            x, y, w, h = cv2.boundingRect(countor)
            cur_radius2 = w**2 + h**2
            if cur_radius2 > radius2:
                radius2 = cur_radius2
                bbox = [x, y, x+w, h+y]
                code = part['label']
        
        if Path(json_path).parent.name not in ["TI01", "TI02"]:
            code = Path(json_path).parent.name
        return (code, bbox)
    except:
        return (None, None)


# def image_crop_by_json(img_path, json_path, method):
#     code, original_bbox = get_original_bbox(json_path)
#     if code is None:
#         return None, None
    
#     bbox = crop_image(img_path=img_path, original_bbox=original_bbox, method=method)
#     return code, bbox

def label_path2img_path(label_path):
    img_path = str(Path(label_path).with_suffix(".jpg"))
    if os.path.exists(img_path):
        return img_path
    return 'INVALID'

def save_crop_img(img_path, box, code, dst_path):
    img_name = Path(img_path).name 
    abs_dst = os.path.join(dst_path, code, img_name)
    os.makedirs(os.path.join(dst_path, code), exist_ok=True)

    new_x_min, new_y_min, new_x_max, new_y_max = box

    img = cv2.imread(img_path, 1)
    crop_img = img[new_y_min:new_y_max, new_x_min:new_x_max]

    cv2.imwrite(abs_dst, crop_img)

def is_None(x):
    return x is None

def main():

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    # root_p = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_data'
    # ps = []
    # for product in os.listdir(root_p):
    #     # for folder in os.listdir(os.path.join(root_p, product)):
    #     ps.append(os.path.join(root_p, product))
    #     # if not os.path.isdir(os.path.join(root_p, product)):
    #     #     continue
    #     # for folder in os.listdir(os.path.join(root_p, product)):
    #     #     ps.append(os.path.join(root_p, product, folder))
    # print(ps)
    
    # dst_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_xml'
    # ds = []
    # for p in ps:
    #     product = str(Path(p).name)
    #     # folder = str(Path(p).name)
    #     ds.append(os.path.join(dst_root, product))

    # p_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_data'
    # ps = []
    # for product in os.listdir(p_root):
    #       ps.append(os.path.join(p_root, product))

    # ds = None

    # if ds is None:
    #     ds = [x+"_xml" for x in ps]

    # for path, dst_path in zip(ps, ds):

    #     tmp = []
    #     for code in os.listdir(path):
    #         imgs = glob.glob(os.path.join(path, code, "**/*.jpg"), recursive=True)
    #         tmp.extend(imgs)

    #     df = pd.DataFrame()
    #     df['img_path'] = tmp

    #     if is_pandarallel:
    #         df.parallel_apply(lambda x:json2xml(img_path=x['img_path'], dst_path=dst_path), axis=1)
    #     else:
    #         df.apply(lambda x:json2xml(img_path=x['img_path'], dst_path=dst_path), axis=1)
    #     print()


    # #####################################################################################
    
    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_data'
    dst_path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_crop'
    method = '299_scale_crop'

    for product in os.listdir(path):
        if product == '236H': continue
        print('product: ', product)
        # for folder in os.listdir(os.path.join(path, product)):
        # print('folder: ', folder)
        jsons = glob.glob(os.path.join(path, product, "*/*.json"))
        if len(jsons) == 0: continue
        df = pd.DataFrame()
        df['json_path'] = jsons
        df['img_path'] = df['json_path'].apply(lambda x: label_path2img_path(x))
        df = df[df['img_path']!='INVALID'].reset_index(drop=True)
        if len(df) == 0:  continue

        dst_path_cur = os.path.join(dst_path, product)

        if is_pandarallel:
            df["code_box"] = df['json_path'].parallel_apply(lambda x:  get_original_bbox(x))
            df['code'] = df['code_box'].apply(lambda x: x[0])
            df['box'] = df['code_box'].apply(lambda x: x[1])
            df = df[~df['box'].apply(lambda x: is_None(x))].reset_index(drop=True)
            df['new_box'] = df.parallel_apply(lambda x: crop_image(img_path=x['img_path'], original_bbox=x['box'], method=method) , axis=1)
            df.parallel_apply(lambda x: save_crop_img(img_path=x['img_path'], box=x['new_box'], code=x['code'], dst_path=dst_path_cur), axis=1)
        else:
            df["code_box"] = df['json_path'].parallelapply_apply(lambda x:  get_original_bbox(x))
            df['code'] = df['code_box'].apply(lambda x: x[0])
            df['box'] = df['code_box'].apply(lambda x: x[1])
            df = df[~df['box'].apply(lambda x: is_None(x))].reset_index(drop=True)
            df['new_box'] = df.apply(lambda x: crop_image(img_path=x['img_path'], original_bbox=x['box'], method=method) , axis=1)
            df.apply(lambda x: save_crop_img(img_path=x['img_path'], box=x['new_box'], code=x['code'], dst_path=dst_path_cur), axis=1)
        print()


if __name__=='__main__':
    main()
